import numpy as np
import matplotlib.pyplot as plt

# 读取文件
# xArr (偏移量x0，横坐标x1)
# yArr (纵坐标)
def loadDataSet(fileName):
    dataMat=np.loadtxt(fileName,delimiter=' ',dtype=int)
    print(dataMat.shape[0])
    xArr=np.zeros((dataMat.shape[0],2)).astype(np.int16)
    yArr=np.array([])
    xArr[:,0]=1
    xArr[:,1]=dataMat[:,0]
    yArr=dataMat[:,1]
    return xArr,yArr

def lwlr(testPoint,xArr,yArr,k):
    xMat=np.mat(xArr)
    yMat=np.mat(yArr)
    yMat=np.transpose(yMat)
    # 创建对角矩阵
    weights=np.mat(np.eye(xMat.shape[0]))
    # 计算权重矩阵weights
    for i in range(weights.shape[0]):
        diffMat=testPoint-xMat[i,:]
        weights[i,i]=np.exp(diffMat*(np.transpose(diffMat))/(-2*k**2))
    xTx=np.transpose(xMat)*(weights*xMat)
    # 计算回归系数向量ws----2*1向量(w0,w1)
    ws=np.linalg.inv(xTx)*(np.transpose(xMat)*(weights*yMat))
    # 得到testpoint(x0,x1)的预测结果----x0w0+x1w1=w0+w1*x1
    return testPoint*ws


def lwlrAll(testArr,xArr,yArr,k):
    yHat=np.zeros(np.shape(testArr)[0])
    for i in range(np.shape(testArr)[0]):
        yHat[i]=lwlr(testArr[i],xArr,yArr,k)
    return yHat.astype(np.int16)


if __name__=="__main__":
    xArr, yArr = loadDataSet('./processed_datasets/region3.txt')
    print(xArr)
    print(yArr)
    print(lwlr([1, len(xArr)+1], xArr, yArr, 2.7))
    yHat = lwlrAll(xArr, xArr, yArr, 2.7)
    print(yHat)

    # 画出图形
    # 获取x坐标从小到大的索引值
    sortIndex = np.argsort(xArr[:, 1], axis=0)
    xSort = xArr[sortIndex]  # n*2矩阵
    fig = plt.figure()
    plt.subplot(111)
    plt.plot(xSort[:, 1], yHat[sortIndex])
    plt.scatter(xArr[:, 1], yArr, s=5, c='red')
    plt.show()



